import pickle

import torch
from torch import nn, optim
from torch.nn import LSTM, modules, Embedding
from config import *
from model import EncoderRNN, DecoderRNN
from pre_process import get_embedding_matrix

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
BASE_DIR = './data/'

word2ix = pickle.load(open(BASE_DIR + 'word_index.pkl', 'rb'))
embedding_matrix = get_embedding_matrix(word_index= word2ix)
encoder = EncoderRNN(WORD_DICT_SIZE+1, EMBEDDING_DIM, ENCODER_DIM, embedding_matrix)

decoder = DecoderRNN(164, 50, 600, eos_id=0, sos_id=163, rnn_cell= 'lstm', bidirectional=True)
test_data, test_tag = pickle.load(open(BASE_DIR+TEST_DATA, 'rb'))

encoder_optimizer = optim.SGD(encoder.parameters(), lr=0.01)

batch_data = test_data[:20]

output, hidden = encoder(torch.tensor(batch_data), device=device)

decoder(encoder_hidden=hidden, encoder_outputs=output)